import os
from fastapi import APIRouter, HTTPException, Header
from pydantic import BaseModel
from database import add_user, get_user_by_username
from datetime import datetime, timedelta
import hashlib
import jwt

user_router = APIRouter()

JWT_SECRET = os.environ.get("JWT_SECRET", "default_secret")
JWT_ALGORITHM = "HS256"
JWT_EXPIRE_MINUTES = 7 * 24 * 60  # 7天

class UserRegister(BaseModel):
    username: str
    password: str
    confirm: str

class UserLogin(BaseModel):
    username: str
    password: str

class UserResetPassword(BaseModel):
    username: str
    new_password: str
    confirm: str

def create_jwt_token(username):
    payload = {
        "username": username,
        "exp": datetime.utcnow() + timedelta(minutes=JWT_EXPIRE_MINUTES)
    }
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)

def verify_jwt_token(token):
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        return payload["username"]
    except Exception:
        return None

def get_current_user(token: str = Header(...)):
    username = verify_jwt_token(token)
    if not username:
        raise HTTPException(status_code=401, detail="无效或过期的token")
    return username

@user_router.post("/user/register")
async def user_register(user: UserRegister):
    if not user.username or not user.password or not user.confirm:
        raise HTTPException(status_code=400, detail="用户名和密码不能为空")
    if user.password != user.confirm:
        raise HTTPException(status_code=400, detail="两次密码输入不一致")
    if not user.password.isdigit() or len(user.password) != 6:
        raise HTTPException(status_code=400, detail="密码必须为6位数字")
    if get_user_by_username(user.username):
        raise HTTPException(status_code=400, detail="用户名已存在")
    password_md5 = hashlib.md5(user.password.encode('utf-8')).hexdigest()
    add_user(user.username, password_md5)
    return {"status": "success"}

@user_router.post("/user/login")
async def user_login(user: UserLogin):
    db_user = get_user_by_username(user.username)
    if not db_user:
        raise HTTPException(status_code=400, detail="用户不存在")
    password_md5 = hashlib.md5(user.password.encode('utf-8')).hexdigest()
    if db_user['password_md5'] != password_md5:
        raise HTTPException(status_code=400, detail="密码错误")
    token = create_jwt_token(user.username)
    return {"status": "success", "token": token, "username": user.username}

@user_router.post("/user/reset_password")
async def reset_password(data: UserResetPassword):
    if not data.username or not data.new_password or not data.confirm:
        raise HTTPException(status_code=400, detail="用户名和新密码不能为空")
    if data.new_password != data.confirm:
        raise HTTPException(status_code=400, detail="两次新密码输入不一致")
    if not data.new_password.isdigit() or len(data.new_password) != 6:
        raise HTTPException(status_code=400, detail="新密码必须为6位数字")
    user = get_user_by_username(data.username)
    if not user:
        raise HTTPException(status_code=400, detail="用户不存在")
    password_md5 = hashlib.md5(data.new_password.encode('utf-8')).hexdigest()
    # 更新数据库
    import sqlite3
    conn = sqlite3.connect('data/journals.db')
    cursor = conn.cursor()
    cursor.execute("UPDATE users SET password_md5=? WHERE username=?", (password_md5, data.username))
    conn.commit()
    conn.close()
    return {"status": "success"}
